package de.distMLP.communication;

import java.io.IOException;

import org.apache.hama.bsp.BSPPeer;
import org.apache.hama.bsp.sync.SyncException;

import de.distMLP.messages.DoubleMessage;
import de.distMLP.messages.HamaMessage;
import de.distMLP.messages.VectorMessage;
import de.distMLP.messages.WeightMatrixMessage;
import de.mlp_distributed.mlp.math.mahout.Matrix;
import de.mlp_distributed.mlp.math.mahout.Vector;

public class SimpleMessageDistributor extends MessageDistributor {

	public SimpleMessageDistributor(final BSPPeer<?, ?, ?, ?, HamaMessage> peer, final String masterPeer) {
		super(peer, masterPeer);
	}

	@Override
	public HamaMessage distributeMessageToAllPeers(final String senderName, final HamaMessage msg, final String[] peerNames)
			throws IOException, SyncException, InterruptedException {
		final BSPPeer<?, ?, ?, ?, HamaMessage> peer = this.getPeer();

		if (peer.getPeerName().equals(senderName)) {
			for (final String peerName : peerNames) {
				peer.send(peerName, msg);
			}
		}

		peer.sync();

		if (peer.getNumCurrentMessages() == 1) {
			return peer.getCurrentMessage();
		}
		return null;
	}

	@Override
	public HamaMessage distributeAndAggegateMessageToMasterPeer(final HamaMessage msg, final String masterPeer) throws IOException,
			SyncException, InterruptedException {
		final BSPPeer<?, ?, ?, ?, HamaMessage> peer = this.getPeer();

		peer.send(masterPeer, msg);

		peer.sync();

		if (peer.getPeerName().equals(this.getMasterPeer()) && (peer.getNumCurrentMessages() > 0)) {
			final double numberMessages = peer.getNumCurrentMessages();
			final HamaMessage firstMessage = peer.getCurrentMessage();

			if (firstMessage instanceof WeightMatrixMessage) {
				return this.handleWeightMatrixMessage(peer, numberMessages, (WeightMatrixMessage) firstMessage);
			} else if (firstMessage instanceof VectorMessage) {
				return this.handleVectorMessage(peer, numberMessages, firstMessage);
			} else if (firstMessage instanceof DoubleMessage) {
				return this.handleDoubleMessage(peer, numberMessages, firstMessage);
			} else {
				throw new UnsupportedOperationException("Message type is not supported!");
			}
		}
		return null;
	}

	public HamaMessage handleDoubleMessage(final BSPPeer<?, ?, ?, ?, HamaMessage> peer, final double numberMessages, final HamaMessage tmp)
			throws IOException {
		double sum = ((DoubleMessage) tmp).getDoubleValue();
		while (peer.getNumCurrentMessages() > 0) {
			final DoubleMessage tmpMessage = (DoubleMessage) peer.getCurrentMessage();
			sum = MessageAggregator.addDouble(sum, tmpMessage.getDoubleValue());
		}
		return new DoubleMessage(sum);
	}

	public HamaMessage handleVectorMessage(final BSPPeer<?, ?, ?, ?, HamaMessage> peer, final double numberMessages, final HamaMessage tmp)
			throws IOException {
		Vector vectorSum = ((VectorMessage) tmp).getVector();
		while (peer.getNumCurrentMessages() > 0) {
			final VectorMessage tmpMessage = (VectorMessage) peer.getCurrentMessage();
			vectorSum = MessageAggregator.addVector(vectorSum, tmpMessage.getVector());
		}
		return new VectorMessage(vectorSum);
	}

	public HamaMessage handleWeightMatrixMessage(final BSPPeer<?, ?, ?, ?, HamaMessage> peer, final double numberMessages,
			final WeightMatrixMessage m) throws IOException {
		boolean stopTraining = false;
		double cost = 0;
		Matrix[] matrixSum = null;

		if (m.isStopTraining()) {
			stopTraining = true;
		}
		matrixSum = m.getMatrix();
		cost += m.getCost();

		while (peer.getNumCurrentMessages() > 0) {
			final WeightMatrixMessage tmpMessage = (WeightMatrixMessage) peer.getCurrentMessage();
			if (tmpMessage.isStopTraining()) {
				stopTraining = true;
			}
			cost += tmpMessage.getCost();

			if (matrixSum == null) {
				matrixSum = tmpMessage.getMatrix();
				continue;
			}
			matrixSum = MessageAggregator.addMatrixArray(matrixSum, tmpMessage.getMatrix());
		}
		return new WeightMatrixMessage(matrixSum, stopTraining, cost);
	}

}
